package com.ada.spark.partitioner

import org.apache.spark.rdd.RDD
import org.apache.spark.{SparkConf, SparkContext}

/**
  * 要实现自定义的分区器，你需要继承 org.apache.spark.Partitioner 类并实现下面三个方法。
  *
  * （1）numPartitions: Int:返回创建出来的分区数。
  *
  * （2）getPartition(key: Any): Int:返回给定键的分区编号(0到numPartitions-1)。
  *
  * （3）equals():Java 判断相等性的标准方法。这个方法的实现非常重要，Spark 需要用这个方法来检查你的分区器对象是否和其他分区器实例相同，这样 Spark 才可以判断两个 RDD 的分区方式是否相同。
  *
  * 需求：将相同后缀的数据写入相同的文件，通过将相同后缀的数据分区到相同的分区并保存输出来实现。
  */
object CustomerPartitioner {
    def main(args: Array[String]): Unit = {
        //初始化配置信息及SparkContext
        val sparkConf: SparkConf = new SparkConf().setAppName("CustomerPartitioner").setMaster("local[*]")
        val sc = new SparkContext(sparkConf)

        val data = sc.parallelize(Array((1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)))

        val par = data.partitionBy(new CustomerPartitioner(2))
        val result: RDD[(Int, (Int, Int))] = par.mapPartitionsWithIndex((index, items) => items.map((index, _)))

        println(result.collect().mkString(","));
        //(0,(2,2)),(0,(4,4)),(0,(6,6)),(1,(1,1)),(1,(3,3)),(1,(5,5))
    }

}


// Entering paste mode (ctrl-D to finish)
class CustomerPartitioner(numParts: Int) extends org.apache.spark.Partitioner {

    //覆盖分区数
    override def numPartitions: Int = numParts

    //覆盖分区号获取函数
    override def getPartition(key: Any): Int = {
        val ckey: String = key.toString
        ckey.substring(ckey.length - 1).toInt % numParts
    }
}

// Exiting paste mode, now interpreting.
//defined class CustomerPartitioner